from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import hashes,serialization
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives.asymmetric import rsa,padding
import os
import json
import base64
from cryptography.hazmat.primitives import padding
from cryptography.hazmat.primitives.asymmetric import padding as asymmetric_padding
from cryptography.hazmat.primitives.asymmetric import padding as asymmetric_padding

global pri_key_or_psd 

# from file_config import pri_key_or_psd
import file_config
# ---------------------------------------------------------------AES256加密--------------------------------------------------------------
# 密钥生成
def generate_key_aes(password, salt):
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32, #密钥长度
        salt=salt,
        iterations=100000,  #密钥迭代次数
        backend=default_backend()
    )
    key = kdf.derive(password)  
    return key  # 返回key值

# 进行加密
def encrypt_aes(message, key):
    salt = os.urandom(16)
    iv = os.urandom(16)
    cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
    encryptor = cipher.encryptor()
    padder = padding.PKCS7(128).padder()
    padded_data = padder.update(message) + padder.finalize()
    ciphertext = encryptor.update(padded_data) + encryptor.finalize()  # 对消息进行加密
    return salt + iv + ciphertext  # 返回salt，iv和ciphertext方便之后切割和解密

# 进行解密
def decrypt_aes(ciphertext, key):
    salt = ciphertext[:16]
    iv = ciphertext[16:32]
    encrypted_data = ciphertext[32:]
    cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
    decryptor = cipher.decryptor()
    decrypted_data = decryptor.update(encrypted_data) + decryptor.finalize()
    unpadder = padding.PKCS7(128).unpadder()
    unpadded_data = unpadder.update(decrypted_data) + unpadder.finalize()
    return unpadded_data



# ---------------------------------------------------------------RSA加密--------------------------------------------------------------

def generate_key_rsa():
    private_key = rsa.generate_private_key(
        public_exponent=65537,
        key_size=2048,
        backend=default_backend()
    )
    public_key = private_key.public_key()
    return private_key, public_key

def encrypt_rsa(message, public_key):
    ciphertext = public_key.encrypt(
        message,
        asymmetric_padding.OAEP(
            mgf=asymmetric_padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None
        )
    )
    return ciphertext

def decrypt_rsa(ciphertext, private_key):
    plaintext = private_key.decrypt(
        ciphertext,
        asymmetric_padding.OAEP(
            mgf=asymmetric_padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None
        )
    )
    return plaintext


# ---------------------------------------------------------------密码处理---------------------------------------------------------------


def encrypt_message_dict(message_dict, pri_key_or_psd, user_nums, salt=None):
    # 确定加密类型
    encrypt_type = message_dict['type']
    # 将字典转换为 JSON 格式的字符串
    message_json = json.dumps(message_dict).encode('utf-8')
    # 调用加密函数
    if salt is None:
        salt = os.urandom(16)
    if encrypt_type == 'aes':  # 使用AES进行加密
        key = generate_key_aes(pri_key_or_psd, salt)
        ciphertext = encrypt_aes(message_json, key)  # 调用 encrypt_aes 函数
    else:   # 使用RSA进行加密
        pri_key_or_psd, pub_key = generate_key_rsa()
        ciphertext = encrypt_rsa(message_json, pub_key)
    # 保存加密数据到文件
    save_encrypted_data(f'file/user{user_nums}/code{user_nums}.txt', salt + ciphertext, encrypt_type, pri_key_or_psd)
    print("已经写入到文件中")
    return salt + ciphertext, encrypt_type, pri_key_or_psd


def decrypt_message_dict(encrypted_data, pri_key_or_psd, type):
    # 进行解密
    salt = encrypted_data[:16]
    ciphertext = encrypted_data[16:]
    if type == 'aes':
        key = generate_key_aes(pri_key_or_psd, salt)
        decrypted_message_json = decrypt_aes(ciphertext, key)
    else:
        decrypted_message_json = decrypt_rsa(ciphertext, pri_key_or_psd)
    # 将 JSON 格式的字符串转换为字典
    return json.loads(decrypted_message_json.decode('utf-8'))


# ---------------------------------------------------------------文件操作---------------------------------------------------------------


def save_encrypted_data(filename, salt_ciphertext, encrypt_type, pri_key_or_psd):
    with open(filename, 'ab') as f:
        # 将 encrypt_type 和 salt_ciphertext 转换为 Base64 编码
        encrypt_type_bytes = base64.b64encode(encrypt_type.encode('utf-8'))
        salt_ciphertext_b64 = base64.b64encode(salt_ciphertext)
        
        # 如果是 RSA 密钥，则进行序列化并进行 Base64 编码
        if encrypt_type == 'rsa':
            pri_key_or_psd = pri_key_or_psd.private_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PrivateFormat.TraditionalOpenSSL,
                encryption_algorithm=serialization.NoEncryption()
            )
        pri_key_or_psd_b64 = base64.b64encode(pri_key_or_psd)
        
        # 将各部分连接起来，中间以空格分隔
        data_to_save = salt_ciphertext_b64 + b' ' + encrypt_type_bytes + b' ' + pri_key_or_psd_b64 + b'\n'
        # 写入文件
        f.write(data_to_save)


def read_encrypted_data(filename, line_num):
    with open(filename, 'rb') as f:
        lines = f.readlines()
        if line_num >= len(lines):
            raise IndexError("Line number out of range")
        
        line = lines[line_num]
        parts = line.split(b' ')
        salt_ciphertext = base64.b64decode(parts[0])
        encrypt_type = base64.b64decode(parts[1]).decode('utf-8')
        
        # 如果是 RSA 密钥，则进行反序列化
        pri_key_or_psd = base64.b64decode(parts[2])
        if encrypt_type == 'rsa':
            pri_key_or_psd = serialization.load_pem_private_key(
                pri_key_or_psd,
                password=None,
                backend=default_backend()
            )
        
        return salt_ciphertext, encrypt_type, pri_key_or_psd

def count_encrypted_data(filename):
    with open(filename, 'rb') as f:
        return len(f.readlines())



# if __name__ == '__main__':
#     # # 示例用法
#     # password = b"your_password"  #传入以字节字符串
#     # message = "Hello, World!"  # 普通字符串
#     # message = message.encode('utf-8')  # 转换为字节字符串，使用UTF-8编码

#     # key = generate_key_aes(password, os.urandom(16))

#     # # 加密
#     # encrypted_data = encrypt(message, key)
#     # save_to_file('file1.txt', [encrypted_data])
#     # print("Encrypted data saved to file1.txt")

#     # # 解密
#     # encrypted_data_from_file = read_from_file('file1.txt')[0]
#     # decrypted_data = decrypt(encrypted_data_from_file, key)
#     # save_to_file('file2.txt', [decrypted_data])
#     # print("Decrypted data saved to file2.txt")
#     message = {"type":"aes","name": "Example", "url": "http://example.com", "username": "user1", "password": "password1"}


#     # 加密
#     encrypted_message,encrypt_type,pri_key_or_psd = encrypt_message_dict(message, pri_key_or_psd)
#     print("Encrypted message:", encrypted_message)

#     # 解密
#     decrypted_message = decrypt_message_dict(encrypted_message, pri_key_or_psd, encrypt_type)
#     print("Decrypted message:", decrypted_message)

# if __name__ == '__main__':
#     pri_key_or_psd = b"your_password"

#     message1 = {"type":"aes","name": "韩泽斌", "url": "http://example.com", "username": "user1", "password": "password1"}
#     message2 = {"type":"rsa","name": "李哲辉", "url": "http://example.com", "username": "user1", "password": "password1"}
#     message3 = {"type":"rsa","name": "王刀郎", "url": "http://example.com", "username": "user1", "password": "password1"}

#     # 加密
#     for i, message in enumerate([message1, message2, message3], 1):
#         encrypted_message, encrypt_type, pri_key_or_psd = encrypt_message_dict(message, pri_key_or_psd)
        
#         # 读取加密数据从文件
#         salt_ciphertext, encrypt_type, pri_key_or_psd = read_encrypted_data('code.txt')

#         # 解密数据
#         decrypted_message = decrypt_message_dict(salt_ciphertext, pri_key_or_psd, encrypt_type)
#         print(f"Decrypted message {i}:", decrypted_message)